//=========================================================================
// Copyright (C) 2024 The C++ Component Model(COMO) Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//=========================================================================

#ifndef __TlsfStdString_H__
#define __TlsfStdString_H__

#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <string>
#include <vector>
#include "tlsf.h"

namespace como {

#ifndef DEFAULT_POOL_SIZE
  #define DEFAULT_POOL_SIZE    10 * 1024 * 1024       // 10MB
#endif

using MemoryHolderOnFailture = void(*)(void);

struct MemoryHolder {
    MemoryHolder(size_t pool_size, void *mem)
        : mUsagePercentage(100)
    {
        mTlsf_memory_pool = mem;
        mTlsf_handle = tlsf_create_with_pool(mTlsf_memory_pool, pool_size);
        if (nullptr == mTlsf_handle) {
            if (nullptr != mOnFailture) {
                mOnFailture();
            }
        }

        mTlsf_pool_size = pool_size;
        this->mUpstream = nullptr;
    }

    MemoryHolder(size_t pool_size, MemoryHolder *upstream = nullptr)
        : mUsagePercentage(100)
    {
        mTlsf_memory_pool = std::malloc(pool_size);
        mTlsf_handle = tlsf_create_with_pool(mTlsf_memory_pool, pool_size);
        if (nullptr == mTlsf_handle) {
            if (nullptr != mOnFailture) {
                mOnFailture();
            }
        }

        mTlsf_pool_size = pool_size;
        this->mUpstream = upstream;
    }

    MemoryHolder(size_t pool_size, void *mem, MemoryHolder *upstream, int32_t usagePercentage)
        : mUsagePercentage(usagePercentage)
    {
        mTlsf_memory_pool = mem;
        mTlsf_handle = tlsf_create_with_pool(mTlsf_memory_pool, pool_size);
        if (nullptr == mTlsf_handle) {
            if (nullptr != mOnFailture) {
                mOnFailture();
            }
        }

        mTlsf_pool_size = pool_size;
        this->mUpstream = upstream;
    }

    ~MemoryHolder() {
        mTlsf_pool_size = 0;
        tlsf_destroy(mTlsf_handle);
        std::free(mTlsf_memory_pool);
    }

    // disable assignment in method copy
    MemoryHolder& operator=(const MemoryHolder&) = delete;

    // allow assignment in method move
    MemoryHolder& operator=(MemoryHolder&& other) = default;

    void *allocate(std::size_t n, std::size_t align) {
        if (nullptr == mTlsf_handle) {
            if (nullptr != mOnFailture) {
                mOnFailture();
            }

            return nullptr;
        }

        void *ptr = tlsf_memalign(mTlsf_handle, align, n);
        if (nullptr == ptr) {
            if (nullptr != mUpstream) {
                return mUpstream->allocate(n, align);
            }

            return nullptr;
        }

        return ptr;
    }

    void deallocate(void* p) {
        if (nullptr == mTlsf_handle) {
            if (nullptr != mOnFailture) {
                mOnFailture();
            }

            return;
        }

        size_t start_pos = reinterpret_cast<size_t>(mTlsf_memory_pool);
        size_t end_pos = start_pos + mTlsf_pool_size;
        if ((reinterpret_cast<size_t>(p) >= start_pos) &&
                                     (reinterpret_cast<size_t>(p) <= end_pos)) {
            tlsf_free(mTlsf_handle, p);
        }
        else if (nullptr != mUpstream) {
            mUpstream->deallocate(p);
        }
        else {
            //throw exception;
            if (nullptr != mOnFailture) {
                mOnFailture();
            }
        }
    }

    void Set_OnFailture(MemoryHolderOnFailture onFailture) {
        mOnFailture = onFailture;
    }

    void Set_Upstream(MemoryHolder *upstream) {
        mUpstream = upstream;
    }

    void         *mTlsf_memory_pool;
    tlsf_t        mTlsf_handle;
    size_t        mTlsf_pool_size;
    int32_t       mUsagePercentage;
    MemoryHolder *mUpstream = nullptr;
    MemoryHolderOnFailture mOnFailture = nullptr;

    static MemoryHolder default_memory_holder;
    static MemoryHolder *como_current_memory_holder;
};

inline MemoryHolder *get_current_memory_holder() {
    return MemoryHolder::como_current_memory_holder;
}

inline void set_current_memory_holder(MemoryHolder* new_memory_holder) {
    MemoryHolder::como_current_memory_holder = new_memory_holder;
}

// Custom allocator using TLSF
template <typename T>
class TlsfAllocator {
public:
    using value_type = T;
    using pointer = T *;
    using const_pointer = const T *;

    using void_pointer = void *;
    using const_void_pointer = const void *;

    using size_type = size_t;
    using difference = std::ptrdiff_t;

    virtual  MemoryHolder *get_default_memory_holder() {
        return MemoryHolder::como_current_memory_holder;
    }

    TlsfAllocator() noexcept {
        mMemory_pool = get_default_memory_holder();
    }

    TlsfAllocator(MemoryHolder *memory_pool) {
        mMemory_pool = memory_pool;
    }

    TlsfAllocator(const TlsfAllocator& other) {
        mMemory_pool = other.mMemory_pool;
    }

    ~TlsfAllocator() noexcept {
        // Do not free the pool in the destructor, as it might be shared
    }

    template <typename U>
    TlsfAllocator(const TlsfAllocator<U>& other) noexcept {
        mMemory_pool = other.mMemory_pool;
    }

    T *allocate(std::size_t n) {
        return static_cast<T*>(mMemory_pool->allocate(n, alignof(T)));
    }

    void deallocate(T* p, std::size_t) noexcept {
        // Free memory using TLSF
        mMemory_pool -> deallocate(p);
    }

    template <typename U>
    struct rebind {
        using other = TlsfAllocator<U>;
    };

    bool operator==(const TlsfAllocator& other) const noexcept { return this->mMemory_pool == other.mMemory_pool; }
    bool operator!=(const TlsfAllocator& other) const noexcept { return this->mMemory_pool != other.mMemory_pool; }

public:
    MemoryHolder *mMemory_pool;
};


//#ifdef COMO_FUNCTION_SAFETY_RTOS
#if 0
using TlsfString = std::basic_string<char, std::char_traits<char>, TlsfAllocator<char>>;

typedef struct tagTlsfYoungAge {
    char c;

    bool operator==(const tagTlsfYoungAge& other) const {
        return c == other.c;
    }

    bool operator<(const tagTlsfYoungAge& other) const {
        return c < other.c;
    }
} TlsfYoungAge;

typedef struct tagTlsfMiddleAge {
    char c;

    bool operator==(const tagTlsfMiddleAge& other) const {
        return c == other.c;
    }

    bool operator<(const tagTlsfMiddleAge& other) const {
        return c < other.c;
    }
} TlsfMiddleAge;

typedef struct tagTlsfOldAge {
    char c;

    bool operator==(const tagTlsfOldAge& other) const {
        return c == other.c;
    }

    bool operator<(const tagTlsfOldAge& other) const {
        return c < other.c;
    }
} TlsfOldAge;

static_assert(sizeof(TlsfYoungAge) == 1, "Error: struct TlsfYoungAge size is not 1!");
static_assert(sizeof(TlsfMiddleAge) == 1, "Error: struct TlsfMiddleAge size is not 1!");
static_assert(sizeof(TlsfOldAge) == 1, "Error: struct TlsfOldAge size is not 1!");

using TlsfYoungAgeString = std::basic_string<TlsfYoungAge,
                    std::char_traits<TlsfYoungAge>, TlsfAllocator<TlsfYoungAge>>;
using TlsfMiddleAgeString = std::basic_string<TlsfMiddleAge,
                    std::char_traits<TlsfMiddleAge>, TlsfAllocator<TlsfMiddleAge>>;
using TlsfOldAgeString = std::basic_string<TlsfOldAge,
                    std::char_traits<TlsfOldAge>, TlsfAllocator<TlsfOldAge>>;
#else
#define TlsfYoungAgeString_IS_std_string    1
#define TlsfMiddleAgeString_IS_std_string   1
#define TlsfOldAgeString_IS_std_string      1

using TlsfYoungAge = char;
using TlsfMiddleAge = char;
using TlsfOldAge = char;
using TlsfString = std::string;
using TlsfYoungAgeString = std::string;
using TlsfMiddleAgeString = std::string;
using TlsfOldAgeString = std::string;
#endif

} // namespace como

#endif // __TlsfAllocator_H__
